#!/usr/bin/env python3
# -*- coding:utf-8 -*-


from PyQt5.QtWidgets import (QMainWindow, QDialog, QWidget,
                             QApplication, QHBoxLayout,
                             QVBoxLayout, QPushButton,
                             QLabel, QGridLayout,
                             QProgressBar, QFileDialog,
                             QLineEdit)
from PyQt5.QtGui import (QPixmap, QImage)
from os import path
import cv2
import numpy as np


class YoltWindow(QMainWindow):
    def __init__(self):
        super().__init__()

        # Yolt 相关参数
        self.use_tiny_yolt = False
        self.weights_file_path = ''
        self.curr_process_image = None
        self.shown_image_w_h = (640, 480)
        self.shown_qimage = QImage()

        # 初始化 GUI
        self.initUI()

    def initUI(self):
        # GUI 相关参数
        self.btn_net_architecture = QPushButton('Use Tiny Yolt')
        self.btn_net_architecture.setCheckable(True)
        self.btn_net_weights = QPushButton('Network Weights')
        self.btn_input_image = QPushButton('Input Image')
        self.lbl_net_architecture = QLineEdit('Fully Yolt')
        self.lbl_net_weights = QLineEdit('Network Weights')
        self.lbl_input_image = QLineEdit('Input Image')
        self.lbl_net_architecture.setReadOnly(True)
        self.lbl_net_weights.setReadOnly(True)
        self.lbl_input_image.setReadOnly(True)
        self.btn_exit = QPushButton('Exit')
        self.process_bar = QProgressBar()
        self.lbl_show_result = QLabel()
        self.pixmap = QPixmap()
        self.lbl_show_result.setPixmap(self.pixmap)

        central_wideget = QWidget()
        self.setCentralWidget(central_wideget)

        # Layout
        main_layout = QGridLayout()
        main_layout.setSpacing(10)
        main_layout.addWidget(self.btn_net_architecture, 0, 0)
        main_layout.addWidget(self.lbl_net_architecture, 0, 1)
        main_layout.addWidget(self.btn_net_weights, 1, 0)
        main_layout.addWidget(self.lbl_net_weights, 1, 1)
        main_layout.addWidget(self.btn_input_image, 2, 0)
        main_layout.addWidget(self.lbl_input_image, 2, 1)
        main_layout.addWidget(self.btn_exit, 3, 0)
        main_layout.addWidget(self.process_bar, 0, 2)
        main_layout.addWidget(self.lbl_show_result, 1, 2)
        main_layout.setColumnStretch(0, 1)
        main_layout.setColumnStretch(1, 1)
        main_layout.setColumnStretch(2, 2)
        central_wideget.setLayout(main_layout)

        # Connect signal and slot
        self.btn_net_architecture.clicked.connect(self.select_yolt_version)
        self.btn_net_weights.clicked.connect(self.select_weights_file)
        self.btn_input_image.clicked.connect(self.select_image)
        self.btn_exit.clicked.connect(self.close)

        self.setWindowTitle("Yolt")
        self.setGeometry(300, 300, 640, 480)
        self.show()

    def select_image(self):
        def numpy_image_to_qimage(raw_image):
            raw_h, raw_w = raw_image.shape[:2]
            if raw_h > self.shown_image_w_h[0] or raw_w > self.shown_image_w_h[1]:
                factor = min(
                    self.shown_image_w_h[0] / raw_w, self.shown_image_w_h[1] / raw_h)
                new_w = round(raw_w * factor)
                new_h = round(raw_h * factor)
                resized_image = cv2.resize(raw_image, (new_w, new_h))
            else:
                resized_image = raw_image

            is_color_image = len(resized_image.shape) == 3
            if is_color_image:
                resized_image = cv2.cvtColor(
                    resized_image, cv2.COLOR_BGR2RGB)
                shown_image_format = QImage.Format_RGB888
            else:
                shown_image_format = QImage.Format_Grayscale8

            h, w = resized_image.shape[:2]
            c = 3 if is_color_image else 1
            new_image = QImage(resized_image.data, w,
                               h, w*c, shown_image_format)
            return new_image

        file_path, _ = QFileDialog.getOpenFileName(
            self, 'Choose An Image as Net Input', '/home/pi', 'Image(*.png *.jpg *.jpeg)')

        if len(file_path) <= 0:
            return None

        file_name = path.basename(file_path)
        self.lbl_input_image.setText(file_name)
        self.curr_process_image = cv2.imread(file_path)
        # self.shown_qimage = numpy_image_to_qimage(self.curr_process_image)
        self.lbl_show_result.setPixmap(
            QPixmap(numpy_image_to_qimage(self.curr_process_image)))

    def select_yolt_version(self, isChecked):
        self.use_tiny_yolt = isChecked
        if isChecked:
            self.lbl_net_architecture.setText('Tiny Yolt')
        else:
            self.lbl_net_architecture.setText('Fully Yolt')

    def select_weights_file(self):
        file_path, _ = QFileDialog.getOpenFileName(
            self, dialog_title, './', file_filter)

        file_name = path.basename(file_path)
        curr_label.setText(file_name)


from sys import argv

if __name__ == "__main__":
    app = QApplication(argv)
    yoltWindow = YoltWindow()
    yoltWindow.show()
    app.exec_()
